import asyncio
import datetime
from typing import List

from llama_index.core.agent.workflow import  FunctionAgent
from llama_index.core.base.llms.types import ChatMessage
from llama_index.core.extractors import KeywordExtractor
from llama_index.core.storage.chat_store.sql import SQLAlchemyChatStore
from llama_index.core.tools import QueryEngineTool
from llama_index.core.vector_stores import SimpleVectorStore
from llama_index.core.schema import TextNode, NodeWithScore
from llama_index.core import Settings, SimpleKeywordTableIndex, SummaryIndex, get_response_synthesizer
from llama_index.embeddings.zhipuai import ZhipuAIEmbedding
from llama_index.core.graph_stores import SimplePropertyGraphStore
from llama_index.core.schema import Document
from pydantic import BaseModel

embed_model = ZhipuAIEmbedding(
    model="embedding-2",
    api_key="f387f5e4837d4e4bba6d267682a957c9.PmPiTw8qVlsI2Oi5"
    # With the `embedding-3` class
    # of models, you can specify the size
    # of the embeddings you want returned.
    # dimensions=1024
)
Settings.embed_model=embed_model

from llama_index.llms.deepseek import DeepSeek

llm = DeepSeek(model="deepseek-chat", api_key="sk-605e60a1301040759a821b6b677556fb")
Settings.llm = llm

from llama_index.core.base.base_selector import (
    BaseSelector,
    MultiSelection,
    SingleSelection,
    SelectorResult,
)
from llama_index.core.selectors.embedding_selectors import EmbeddingSingleSelector
from llama_index.core.selectors.llm_selectors import (
    LLMMultiSelector,
    LLMSingleSelector,
)
from llama_index.core.selectors.pydantic_selectors import (
    PydanticMultiSelector,
    PydanticSingleSelector,
)

'''
selector=LLMSingleSelector.from_defaults()

# 定义候选选项
choices = [
    "使用向量检索获取文档片段",
    "调用知识图谱查询关系",
    "直接生成自由文本回答"
]

# 执行选择（输入用户问题）
result = selector.select(
    choices,
    query="知识图谱？"
)
print(f"推荐方案: {result}")  # 输出: "使用向量检索获取文档片段"

choices = [
    "使用向量检索获取文档片段",
    "调用知识图谱查询关系",
    "直接生成自由文本回答"
]

selector=EmbeddingSingleSelector.from_defaults()

result =selector.select(choices,'知识图谱')
print(f"推荐方案: {result}")
# 初始化多模态选择器

'''

from typing import Literal
from pydantic import BaseModel, Field
from llama_index.core.selectors import PydanticSingleSelector
from llama_index.core.tools import QueryEngineTool

# 自定义决策维度模型
class SelectionCriteria(BaseModel):
    domain_match: float = Field(..., ge=0, le=1, description="领域匹配度")
    term_frequency: int = Field(..., description="关键词命中次数")
    preferred_tool: Literal["finance", "medical", "legal", "general"]
    fallback_reason: str = Field(..., description="备选工具使用原因")

from llama_index.core.agent import ReActChatFormatter
from llama_index.core.agent.react.output_parser import ReActOutputParser
from llama_index.core.tools import FunctionTool
from llama_index.core.llms import ChatMessage



# 配置增强选择器
selector = PydanticSingleSelector.from_defaults(
)
from llama_index.core.tools import ToolMetadata

tool_choices = [
    ToolMetadata(

        name="covid_nyt",
        description=("This tool contains a NYT news article about COVID-19"),
    ),
    ToolMetadata(
        name="covid_wiki",
        description=("This tool contains the Wikipedia page about COVID-19"),
    ),
    ToolMetadata(
        name="covid_tesla",
        description=("This tool contains the Wikipedia page about apples"),
    ),
]
selector_result = selector.select(
    tool_choices, query="Tell me more about COVID-19"
)

print(selector_result)

